Replace cuda::atomic_ref with plain CUDA atomics#10001
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/ok to test 7608ed3 |
PointKernel
added
cuco
This comment has been minimized.
This comment has been minimized.
|
/ok to test 45c0caa |
This comment has been minimized.
This comment has been minimized.
|
/ok to test 07c9fbc |
🥳 CI Workflow Results🟩 Finished in 45m 21s: Pass: 100%/57 | Total: 13h 55m | Max: 45m 20s | Hits: 82%/46960See results here. |
|
@wmaxey Requesting your review since you're the team's atomic expert. |
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesThis change adds CUCO-scoped CUDA atomic helpers for supported integral widths and replaces selected CUCO plain CUDA atomics
Assessment against linked issues
Suggested reviewers: Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
cudax/include/cuda/experimental/__cuco/detail/utility/atomic.cuh (1)
137-171: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valuesuggestion: Narrow-width
__atomic_compare_exchangecomputes__aligned/__offsetfrom the raw address without asserting_Tp's natural alignment. If a future 1/2-byte_Tpis only 1-byte aligned,__offsetcould put the packed window across a 4-byte boundary, corrupting adjacent memory. Consider a static/runtime assert on alignment, or document the natural-alignment requirement.cudax/test/cuco/utility/test_atomic.cu (1)
38-64: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winsuggestion: Coverage gap:
__atomic_fetch_maxis only exercised for a 4-byte type (line 44). No test covers the 8-byte signed path, which is exactly the combination flagged as needing verification inatomic.cuh's__atomic_fetch_max. See the root-cause comment there.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0c05c2d0-decd-47c8-bf15-79a6884bd28d
📒 Files selected for processing (6)
cudax/include/cuda/experimental/__cuco/detail/hyperloglog/hyperloglog_impl.cuhcudax/include/cuda/experimental/__cuco/detail/open_addressing/kernels.cuhcudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_ref_impl.cuhcudax/include/cuda/experimental/__cuco/detail/utility/atomic.cuhcudax/test/CMakeLists.txtcudax/test/cuco/utility/test_atomic.cu
| const auto key_cas_success = ::cuda::experimental::cuco::detail::__atomic_compare_exchange<_Scope>( | ||
| &__address->first, __expected_key, __key_type{__desired.first}); | ||
| auto payload_cas_success = ::cuda::experimental::cuco::detail::__atomic_compare_exchange<_Scope>( | ||
| &__address->second, __expected_payload, mapped_type{__desired.second}); | ||
|
|
||
| // if __key success | ||
| if (key_cas_success) | ||
| { | ||
| while (!payload_cas_success) | ||
| { | ||
| payload_cas_success = payload_ref.compare_exchange_strong( | ||
| __expected_payload = empty_value_sentinel(), __desired.second, ::cuda::memory_order_relaxed); | ||
| __expected_payload = empty_value_sentinel(); | ||
| payload_cas_success = ::cuda::experimental::cuco::detail::__atomic_compare_exchange<_Scope>( | ||
| &__address->second, __expected_payload, mapped_type{__desired.second}); | ||
| } | ||
| return __insert_result::__success; | ||
| } | ||
| else if (payload_cas_success) | ||
| { | ||
| // This is insert-specific, cannot for `erase` operations | ||
| payload_ref.store(empty_value_sentinel(), ::cuda::memory_order_relaxed); | ||
| ::cuda::experimental::cuco::detail::__atomic_store<_Scope>(&__address->second, empty_value_sentinel()); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
important: Logic transcription from the prior atomic_ref version is correct. The __atomic_store call at line 765 depends on mapped_type being 4/8 bytes per atomic.cuh's current constraints — see the root-cause comment on atomic.cuh's __atomic_load/__atomic_store for the narrow-payload risk this call inherits.
| auto __expected_key = __expected.first; | ||
| const auto success = | ||
| key_ref.compare_exchange_strong(__expected_key, __key_type{__desired.first}, ::cuda::memory_order_relaxed); | ||
| const auto success = ::cuda::experimental::cuco::detail::__atomic_compare_exchange<_Scope>( | ||
| &__address->first, __expected_key, __key_type{__desired.first}); | ||
|
|
||
| // if __key success | ||
| if (success) | ||
| { | ||
| ::cuda::atomic_ref<mapped_type, _Scope> payload_ref(__address->second); | ||
| payload_ref.store(__desired.second, ::cuda::memory_order_relaxed); | ||
| ::cuda::experimental::cuco::detail::__atomic_store<_Scope>(&__address->second, mapped_type{__desired.second}); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
important: Same narrow-payload dependency as back_to_back_cas: the __atomic_store at line 801 requires a 4/8-byte mapped_type. See the root-cause comment on atomic.cuh.
| _Value __current; | ||
| // TODO exponential backoff strategy | ||
| do | ||
| { | ||
| __current = __ref.load(::cuda::std::memory_order_relaxed); | ||
| __current = ::cuda::experimental::cuco::detail::__atomic_load<_Scope>(&__slot); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
important: __atomic_load<_Scope> requires a 4/8-byte _Value; __wait_for_payload is generic over any slot type, so this inherits the same narrow-width gap flagged on atomic.cuh.
| template <::cuda::thread_scope _Scope, class _Tp> | ||
| [[nodiscard]] _CCCL_DEVICE_API _Tp __atomic_load(const _Tp* __address) noexcept | ||
| { | ||
| ::cuda::experimental::cuco::detail::__validate_atomic_type<_Tp>(); | ||
| static_assert(sizeof(_Tp) == 4 || sizeof(_Tp) == 8, "cuCO atomic load requires a 4 or 8 byte type"); | ||
|
|
||
| using __word_type = __atomic_word_t<_Tp>; | ||
| auto* const __word_address = reinterpret_cast<__word_type*>(const_cast<_Tp*>(__address)); | ||
| constexpr __word_type __zero = 0; | ||
| const auto __old = ::cuda::experimental::cuco::detail::__atomic_cas_word<_Scope>(__word_address, __zero, __zero); | ||
| return ::cuda::std::bit_cast<_Tp>(__old); | ||
| } | ||
|
|
||
| template <::cuda::thread_scope _Scope, class _Tp> | ||
| _CCCL_DEVICE_API void __atomic_store(_Tp* __address, _Tp __value) noexcept | ||
| { | ||
| ::cuda::experimental::cuco::detail::__validate_atomic_type<_Tp>(); | ||
| static_assert(sizeof(_Tp) == 4 || sizeof(_Tp) == 8, "cuCO atomic store requires a 4 or 8 byte type"); | ||
|
|
||
| using __word_type = __atomic_word_t<_Tp>; | ||
| auto* const __word_address = reinterpret_cast<__word_type*>(__address); | ||
| const __word_type __desired = ::cuda::std::bit_cast<__word_type>(__value); | ||
| (void) ::cuda::experimental::cuco::detail::__atomic_exchange_word<_Scope>(__word_address, __desired); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
important: __atomic_load/__atomic_store only support 4/8-byte types (static_assert(sizeof(_Tp) == 4 || sizeof(_Tp) == 8, ...)), with no narrow (<=2 byte) fallback, unlike __atomic_compare_exchange which explicitly implements a packed-CAS path for small types. The previous cuda::atomic_ref<T, Scope> supported arbitrary trivially-copyable sizes. open_addressing_ref_impl.cuh's __wait_for_payload, back_to_back_cas, and cas_dependent_write call __atomic_load/__atomic_store generically on payload types that are not guaranteed to be >=4 bytes (e.g. a key/payload pair combined size >8 bytes with a 1-2 byte payload). Such instantiations will fail to compile where the old atomic_ref-based code compiled fine.
🛠️ Suggested direction
Mirror the packed-CAS technique already implemented in __atomic_compare_exchange to add a narrow-width load (peek via CAS(0,0) on the aligned word, masked/shifted) and narrow-width store (CAS loop overwriting only the target bits).
1 participant
Description
Closes #9983
This PR replaces
cuda::atomic_refwith plain CUDA atomics in the cuco open-addressing and HyperLogLog hot paths.It adds internal helpers for the required relaxed compare-and-exchange, load/store, add, and max operations.
Checklist